// File provided for Microsoft Shared Source.
// Copyright (c) Microsoft Corporation. All rights reserved.

#pragma once
#include "shobjidl.h"
#include "guiddef.h"
#include "wmifldr.h"

struct WMIENUMINFO
{
    PCWSTR pszComputer;
    PCWSTR pszNamespace;
    PCWSTR pszClass;
    PCWSTR pszInstance;

    HRESULT GetServices(__out IWbemServices **ppServices);
    HRESULT GetInstancePath(BOOL fAbsolute, __out BSTR *pbstrInstancePath);
    HRESULT GetInstance(__in IWbemServices *pSvc, __out IWbemClassObject **ppInstance);
    HRESULT GetInstance(__out IWbemClassObject **ppInstance);
    HRESULT InitFromPath(__in PCWSTR pszParsingPath);
    
    WMIENUMINFO();
    ~WMIENUMINFO();

private:
    WCHAR szServer[MAX_PATH];
    WCHAR szNamespace[MAX_PATH];
    WCHAR szClassName[MAX_PATH];
    WCHAR szInstance[MAX_PATH];
};

class CMultiEnumerator : public IEnumWbemClassObject
{
private:
    LONG _cRef;

    size_t _iEnumCurrent;
    CAtlArray< CComPtr<IEnumWbemClassObject> > _srgpEnumerators;

    CMultiEnumerator();
    ~CMultiEnumerator();

    friend class CComPtr<CMultiEnumerator>;
    friend class _NoAddRefReleaseOnCComPtr<CMultiEnumerator>;

public:
    static HRESULT Create(__deref_out CMultiEnumerator **ppMultiEnum);
    HRESULT AddEnumerator(__in IEnumWbemClassObject *pEnumWbemClassObject);

    // IUnknown
    IFACEMETHODIMP QueryInterface(REFIID, void **ppv);
    IFACEMETHODIMP_(ULONG) AddRef()
    {
        return InterlockedIncrement(&_cRef);
    }
    IFACEMETHODIMP_(ULONG) Release()
    {
        LONG cRef = InterlockedDecrement(&_cRef);
        if (0 == cRef)
        {
            delete this;
        }
        return cRef;
    }

    // IEnumWbemClassObject
    IFACEMETHODIMP Reset();
    IFACEMETHODIMP Next(__in long lTimeout, __in ULONG uCount, __out_ecount_part(uCount, *puReturned) IWbemClassObject **apObjects, __out ULONG *puReturned);
    IFACEMETHODIMP NextAsync(__in ULONG uCount, __in_opt IWbemObjectSink *pSink);
    IFACEMETHODIMP Clone(__deref_out IEnumWbemClassObject **ppEnum);
    IFACEMETHODIMP Skip(__in long lTimeout, __in ULONG nCount);
};

class CWMIInstanceEnumerator : public IEnumIDList
{
private:
    LONG _cRef;

    CComPtr<IWbemClassObject> _spInstance;
    CAtlArray<CComBSTR> _asbstrPropertyNames;
    size_t _iProperty;

    HRESULT _Initialize(__in WMIENUMINFO *pEnumInfo);

    CWMIInstanceEnumerator() { _cRef = 1; _iProperty = 0; }
    ~CWMIInstanceEnumerator() {}
public:
    template <class T> static HRESULT Create(T creationData, REFIID riid, VOID **ppv);

    // IUnknown
    IFACEMETHODIMP QueryInterface(REFIID, void **ppv);
    IFACEMETHODIMP_(ULONG) AddRef()
    {
        return InterlockedIncrement(&_cRef);
    }
    IFACEMETHODIMP_(ULONG) Release()
    {
        LONG cRef = InterlockedDecrement(&_cRef);
        if (0 == cRef)
        {
            delete this;
        }
        return cRef;
    }

    // IEnumIDList
    IFACEMETHODIMP Next(ULONG celt, PITEMID_CHILD *rgelt, ULONG *pceltFetched);
    IFACEMETHODIMP Reset();
    IFACEMETHODIMP Clone(IEnumIDList **ppenum);
    IFACEMETHODIMP Skip(ULONG);
};

class CWMIEnumerator :
    public IEnumIDList
{
private:
    LONG _cRef;

    CComPtr<IEnumWbemClassObject> _spEnum;

    CWMIEnumerator(void);
    ~CWMIEnumerator(void);

    HRESULT _Initialize(__in WMIENUMINFO *pEnumInfo);
    HRESULT _Initialize(__in CWMIEnumerator *pWMIEnum);

public:
    template <class T> static HRESULT Create(T creationData, REFIID riid, VOID **ppv);

    // IUnknown
    IFACEMETHODIMP QueryInterface(REFIID, void **ppv);
    IFACEMETHODIMP_(ULONG) AddRef()
    {
        return InterlockedIncrement(&_cRef);
    }
    IFACEMETHODIMP_(ULONG) Release()
    {
        LONG cRef = InterlockedDecrement(&_cRef);
        if (0 == cRef)
        {
            delete this;
        }
        return cRef;
    }

    // IEnumIDList
    IFACEMETHODIMP Next(ULONG celt, PITEMID_CHILD *rgelt, ULONG *pceltFetched);
    IFACEMETHODIMP Reset();
    IFACEMETHODIMP Clone(IEnumIDList **ppenum);
    IFACEMETHODIMP Skip(ULONG);
};

template<class T>
HRESULT CWMIEnumerator::Create(T createData, REFIID riid, VOID **ppv)
{
    HRESULT hr;
    *ppv = NULL;

    CWMIEnumerator *pEnumerator = new CWMIEnumerator();
    hr = pEnumerator ? S_OK : E_OUTOFMEMORY;
    if (SUCCEEDED(hr))
    {
        hr = pEnumerator->_Initialize(createData);
        if (SUCCEEDED(hr))
        {
            hr = pEnumerator->QueryInterface(riid, ppv);
        }
        pEnumerator->Release();
    }

    TRACEFNEXIT();
    return hr;
}

template<class T>
HRESULT CWMIInstanceEnumerator::Create(T createData, REFIID riid, VOID **ppv)
{
    HRESULT hr;
    *ppv = NULL;

    CWMIInstanceEnumerator *pEnumerator = new CWMIInstanceEnumerator();
    hr = pEnumerator ? S_OK : E_OUTOFMEMORY;
    if (SUCCEEDED(hr))
    {
        hr = pEnumerator->_Initialize(createData);
        if (SUCCEEDED(hr))
        {
            hr = pEnumerator->QueryInterface(riid, ppv);
        }
        pEnumerator->Release();
    }

    TRACEFNEXIT();
    return hr;
}